import logging
from datetime import datetime
from collections.abc import Generator
from concurrent.futures import ThreadPoolExecutor
import json
from typing import Dict, Any
from importlib import import_module

from langgraph.graph.state import CompiledStateGraph

from csagent.core.config.schema import AgentConfig 
from csagent.core.utils import default_short_memory

logger = logging.getLogger()

class Agent:
    def __init__(self):
        self.config: AgentConfig = None
        self.graph: CompiledStateGraph = None
        self.redis_cli = None
        self.es_cli = None
        self.memory_executor = ThreadPoolExecutor(max_workers=5)
    
    def _get_short_term_memory(self, user_id, session_id):
        if not user_id:
            logger.warning(f"user_id is empty")
        if not session_id:
            logger.warning(f"session_id is empty")
        key = f"short_mem:{user_id}_{session_id}"
        short_term = self.redis_cli.hgetall(key)
        if not short_term or not isinstance(short_term, dict):
            return {}
        return short_term
    
    def _get_long_term_memory(self, user_id):
        if not user_id:
            logger.warning(f"user_id is empty")
        doc = self.es_cli.search(index="t_user_profile", data={"query":{"term":{"user_id":user_id}}})
        hits = doc.get('hits', {}).get('hits', [])
        if not hits or not isinstance(hits, list):
            return {}
        return hits[0].get('_source', {})

    def _get_memory(self, data):
        user_id = data.get('user_id', '')
        session_id = data.get('session_id', '')
        short_term_config = self.config.memory.short_term
        if short_term_config.read:
            data.setdefault('short_term_memory', {})
            short_term_memory = self._get_short_term_memory(user_id, session_id)
            data['short_term_memory'].update(short_term_memory)
            logger.info("获取短期记忆:%s" % json.dumps(short_term_memory, ensure_ascii=False))
        long_term_config = self.config.memory.long_term
        if long_term_config.read:
            data.setdefault('long_term_memory', {})
            long_term_memory = self._get_long_term_memory(user_id)
            data['long_term_memory'].update(long_term_memory)
            logger.info("获取长期记忆:%s" % json.dumps(long_term_memory, ensure_ascii=False))
    
    def _write_short_memory(self, data, short_term_config):
        log_id = data.get('log_id', '')
        user_id = data.get('user_id', '')
        session_id = data.get('session_id', '')
        strategy = short_term_config.strategy
        try:
            if not strategy:
                strategy_func = default_short_memory
            else:
                module_path, class_name = strategy.rsplit(".", 1)
                module = import_module(module_path)
                strategy_func = getattr(module, class_name)
            short_term_memory = strategy_func(data, short_term_config)
        except Exception as e:
            logger.warning(f"短期记忆生成失败:%s" % json.dumps({'user_id':user_id, 'session_id':session_id, 'log_id':log_id, 'err':str(e)}, ensure_ascii=False))
            return
        key = f"short_mem:{user_id}_{session_id}"
        if not short_term_memory:
            return
        short_term_memory['update_time'] = datetime.now().strftime("%Y-%m-%dT%H:%M:%S")
        try:
            self.redis_cli.hset(key, mapping=short_term_memory)
            if short_term_config.ttl > 0:
                self.redis_cli.expire(key, short_term_config.ttl)
        except Exception as e:
            logger.warning(f"短期记忆更新失败:%s" % json.dumps({'user_id':user_id, 'session_id':session_id, 'log_id':log_id, 'err':str(e)}, ensure_ascii=False))

    def _write_memory(self, data):
        short_term_config = self.config.memory.short_term
        if short_term_config.write and data['code'] == 0:
            self.memory_executor.submit(self._write_short_memory, data, short_term_config)
            #threading.Thread(target=self._write_short_memory, args=(data, short_term_config), daemon=False).start()
    
    def run(self, input_data: Dict[str, Any], stream=False):
        """执行 Agent"""
        # 获取记忆
        self._get_memory(input_data)
        final_res = input_data
        if stream:
            return self.run_stream(input_data)      
        results = list(self.run_stream(input_data))
        if not results:
            input_data['code'] = 2000
            input_data['err_msg'] = "服务异常"
            return input_data
        final_res = results[-1]
        return final_res
    
    def run_stream(self, input_data: Dict[str, Any]) -> Generator[Any]:
        final_res = input_data
        for (key, value) in self.graph.stream(input_data, stream_mode=["custom", "values"]):
            if key == 'values':
                final_res = value
            else:
                yield value
        self._write_memory(final_res)
        yield final_res

